# coding:utf-8
import numpy as np
# 参考网址：https://www.cnblogs.com/qw12/p/5697206.html
# https://blog.csdn.net/u010866505/article/details/77897632

def qq(y,alpha,mu,sigma,K,gama):#计算Q函数
    gsum=[]
    n=len(y)
    for k in range(K):
            gsum.append(np.sum([gama[j,k] for j in range(n)]))
    return np.sum([g*np.log(ak) for g,ak in zip(gsum,alpha)])+\
           np.sum([[np.sum(gama[j,k]*(np.log(1/np.sqrt(2*np.pi))-
                                      np.log(np.sqrt(sigma[k]))-1/2/sigma[k]*(y[j]-mu[k])**2))
                    for j in range(n)] for k in range(K)])  #《统计学习方法》中公式9.29有误


def phi(mu,sigma,y): #计算phi
    return 1/(np.sqrt(2*np.pi*sigma)*np.exp(-(y-mu)**2/2/sigma))


def gama(alpha,mu,sigma,i,k): #计算gama
    sumak=np.sum([[a*phi(m,s,i)] for a,m,s in zip(alpha,mu,sigma)])
    return alpha[k]*phi(mu[k],sigma[k],i)/sumak


def dataN(length,k):#生成数据
    y=[np.random.normal(5*j,j+5,int(length/k)) for j in range(k)]
    return y


def EM(y,K,iter=1000): #EM算法
    n = len(y)
    sigma=[10]*K
    mu=range(K)
    alpha=np.ones(K)
    qqold,qqnew=0,0
    for it in range(iter):
        gama2=np.ones((n,K))
        for k in range(K):
            for i in range(n):
                gama2[i,k]=gama(alpha,mu,sigma,y[i],k)
        for k in range(K):
            sum_gama=np.sum([gama2[j,k] for j in range(n)])
            mu[k]=np.sum([gama2[j, k] * y[j] for j in range(n)])/sum_gama
            sigma[k]=np.sum([gama2[j,k]*(y[j]-mu[k])**2 for j in range(n)])/sum_gama
            alpha[k]=sum_gama/n
        qqnew=qq(y,alpha,mu,sigma,K,gama2)
        if abs(qqold-qqnew)<0.000001:
            break
        qqold=qqnew
    return alpha,mu,sigma

N = 500
k=2
data=dataN(N,k)
y=np.reshape(data,(1,N))
a,b,c = EM(y[0], k)
print(a,b,c)
# iter=180
#[ 0.57217609  0.42782391] [4.1472879054766887, 0.72534713118155769] [44.114682884921415, 24.676116557533351]

sigma = 6  #网上的数据
miu1 = 40
miu2 = 20
X = np.zeros((1, N))
for i in range(N):
    if np.random.random() > 0.5:
        X[0, i] = np.random.randn() * sigma + miu1
    else:
        X[0, i] = np.random.randn() * sigma + miu2
a,b,c = EM(X[0], k)
print(a,b,c)
# iter=114
#[ 0.44935959  0.55064041] [40.561782615819361, 21.444533254494189] [33.374144230703514, 51.459622219329155]